import arrow
import traceback

from pymongo import MongoClient, UpdateOne

from conf.settings import CLOUDMONITOR_MONGODB, DATE_FMT
from utility.utils import rollback, backup
from utility.log import logger

# db = MongoClient(CLOUDMONITOR_MONGODB)['cloudmonitor']
db = MongoClient(host='localhost', port=27017)['cloudmonitor']

now_datetime = arrow.utcnow()
now_datetime = now_datetime.datetime
now = now_datetime.isoformat()


def main():
    try:
        backup(('balances', 'packages', 'charges'), logger, db, now)
        plan = db['packages'].find_one({'key': 'scanv_custom_a'})
        temp_dict = {
            'vulMaxPage': 100000,
            'securityEventMaxPage': 100000,
            'contentMaxPage': 1000,
            'vulMaxDepth': 10,
            'securityEventMaxDepth': 10,
            'contentMaxDepth': 10,
        }

        balances_update_list = []
        charges_update_list = []
        for balances in db['balances'].find({'planId': plan['_id']}):
            balances_update_list.append(UpdateOne({'_id': balances['_id']}, {'$set': {
                'balance.vulMaxPage.total': 100000,
                'balance.securityEventMaxPage.total': 100000,
                'balance.contentMaxPage.total': 1000,
                'balance.vulMaxDepth.total': 10,
                'balance.securityEventMaxDepth.total': 10,
                'balance.contentMaxDepth.total': 10,
            }}))

            for k, v in temp_dict.items():
                query = {'uid': balances['uid'], 'packId': plan['_id'], 'funcType': k, 'chargeType': 'plan'}
                if balances.get('planExpireTime'):
                    charges_update_list.append(UpdateOne(query, {'$set': {f'total': v, 'updateTime': now_datetime, 'isDelete': False, 'expireTime': balances['planExpireTime']}}))
                else:
                    charges_update_list.append(UpdateOne(query, {'$set': {f'total': v, 'updateTime': now_datetime, 'isDelete': False}}))


        if balances_update_list:
            db['balances'].bulk_write(balances_update_list)
        if charges_update_list:
            db['charges'].bulk_write(charges_update_list)

    except Exception as e:
        traceback.print_exc()
        logger.exception(e)
        rollback(('balances', 'packages', 'charges'), logger, db, now)


if __name__ == '__main__':
    main()

